Skip to content

perf: fuse quantized_dense onto mx::quantized_matmul kernel - #200

Merged
ausimian merged 1 commit into
mainfrom
perf/fused-quantized-matmul
Jul 7, 2026
Merged

perf: fuse quantized_dense onto mx::quantized_matmul kernel#200
ausimian merged 1 commit into
mainfrom
perf/fused-quantized-matmul

Conversation

@ausimian

@ausimian ausimian commented Jul 7, 2026

Copy link
Copy Markdown
Owner

Summary

Emily.Quantization.Layers.quantized_dense/4 — the Axon layer backing every quantized projection — now lowers to the fused mx::quantized_matmul kernel (via Emily.Quantization.quantized_matmul_defn/2) instead of dequantizing the full weight to bf16 and running a dense Nx.dot. Inside the compiled decode :while loop the old path re-dequantized the entire model on every token, making quantized inference slower than dense.

Closes #197. This graduates the already-built-but-unwired fused quantized matmul (see PLAN.md M10.5 / M11 — "Option 2 remains available if the two-kernel-vs-fused gap materially hurts a real workload"). The motivating workload is the quantized gap highlighted in elixir-nx/emlx#124.

Results

Same machine, native lane (compiler: Emily.Compiler, native: true, native_fallback: :raise), Qwen/Qwen3-0.6B, 64-token greedy decode:

Lane tok/s vs dense
Dense f32 16.7 1.0×
Quantized 4-bit — before 2.74 0.16× (6× slower)
Quantized 4-bit — after 36.1 2.2× faster
  • ~13× faster end-to-end on quantized decode; quantization flips from slower than dense to faster than dense — the memory-bandwidth win it's supposed to deliver.
  • Per-matmul micro-bench: fused is 2.3–4.6× faster than dequant+dot (largest on the bandwidth-bound MLP projections).
  • Greedy output is byte-identical before/after.

(Absolute tok/s are from a box slower than the emlx PR's M4 Max, so they aren't directly comparable to that PR's numbers; the ratios and the "quantized now beats dense" result are the machine-independent findings.)

Changes

  • lib/emily/quantization/layers.ex — wire the layer to quantized_matmul_defn/2; layout/mode/bits/group_size are read off the %QuantizedWeight{}. Non-Emily backends keep the composed dequantize_defn + Nx.dot fallback via the block's default fun.
  • test/emily/quantization/layers_test.exs — two native-lane tests (native_fallback: :raise) proving quantized_dense genuinely lowers to the fused opcode rather than falling back to op-by-op eval.
  • bench/qmm_microbench.exs — per-matmul before/after (no model download needed).
  • bench/qwen3_quantized_tps.exs — end-to-end quantized Qwen3-0.6B tok/s (run MIX_ENV=test mix run …).

Testing

  • mix precommit clean: compile --warnings-as-errors, format, credo --strict, docs build, and 40 doctests / 79 properties / 737 tests, 0 failures.
  • Existing layers_test already used the eager fused kernel as its oracle, so no tolerance changes were needed.

Emily.Quantization.Layers.quantized_dense/4 dequantized the full weight
to bf16 and ran a dense Nx.dot on every call. Inside the compiled decode
while-loop that re-dequantized the entire model on every token, making
quantized inference slower than dense.

Lower to the fused mx::quantized_matmul kernel (via
Emily.Quantization.quantized_matmul_defn/2) instead, streaming the packed
low-bit weights directly. Layout/mode/bits/group_size are read off the
%QuantizedWeight{}; non-Emily backends keep the composed dequantize +
Nx.dot fallback via the block's default fun.

On 4-bit Qwen3-0.6B native greedy decode this is ~13x faster end-to-end
(2.7 -> 36 tok/s on the dev box), and quantized inference is now faster
than dense rather than slower. Per-matmul it is 2.3-4.6x faster than the
dequantize + Nx.dot path.

Add native-lane tests (native_fallback: :raise) proving quantized_dense
lowers to the fused opcode, plus bench/qmm_microbench.exs (per-matmul)
and bench/qwen3_quantized_tps.exs (end-to-end).

Closes #197.
@ausimian

ausimian commented Jul 7, 2026

Copy link
Copy Markdown
Owner Author

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. 🎉

Reviewed commit: 61b36f4582

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

@ausimian
ausimian merged commit d5a35db into main Jul 7, 2026
3 checks passed
@ausimian
ausimian deleted the perf/fused-quantized-matmul branch July 7, 2026 16:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add a fused quantized_matmul Axon layer (graduate QMM from the notebook)

1 participant